import math
from typing import Tuple, Union, Literal
import pygame
import pygame.gfxdraw
import numpy as np
from .railsim.libapi import DIRECTION, Track

Number = Union[int, float]
SemaphoreStates = Literal["STOP", "PASS", "WARNING"]
_id = 0


def new_id():
    global _id
    _id += 1
    return _id


def rotate_points(points, angle_degrees, a, b):
    """
    Rotate points around a point (a, b) by a certain angle.

    Parameters:
    points: numpy array of shape (n, 2) where n is the number of points
    angle_degrees: The angle in degrees to rotate the points
    a, b: The x and y coordinates of the rotation center

    Returns:
    A numpy array of shape (n, 2) containing the rotated points
    """

    # Convert the angle from degrees to radians
    angle_radians = np.radians(angle_degrees)

    # Create the rotation matrix
    rotation_matrix = np.array(
        [
            [np.cos(angle_radians), -np.sin(angle_radians)],
            [np.sin(angle_radians), np.cos(angle_radians)],
        ]
    )

    # Translate points to the rotation center
    translated_points = points - np.array([a, b])

    # Apply the rotation matrix
    rotated_points = np.dot(translated_points, rotation_matrix.T)

    # Translate points back
    final_points = rotated_points + np.array([a, b])

    return final_points


def draw_equilateral_triangle(surface, color, center, side_length, direction):
    # 计算三角形的高度
    height = (math.sqrt(3) / 2) * side_length
    # 计算底边中点到顶点的距离
    mid_to_top = side_length / 2
    # 计算底边两个顶点的坐标
    top_point = (center[0], center[1] - height / 2)
    left_point = (center[0] - mid_to_top, center[1] + height / 2)
    right_point = (center[0] + mid_to_top, center[1] + height / 2)

    # 根据方向调整三角形的顶点顺序
    # if direction == "left":
    #     points = [top_point, left_point, right_point]
    # elif direction == "right":
    #     points = [top_point, right_point, left_point]

    # 绘制三角形
    pygame.draw.polygon(
        surface,
        color,
        rotate_points(
            [top_point, left_point, right_point],
            90 if direction == "right" else -90,
            center[0],
            center[1],
        ),
    )


class TurnoutJunctionModel:
    def __init__(
        self,
        main_pos_0: Tuple[Number, Number],
        main_pos_1: Tuple[Number, Number],
        side_pos: Tuple[Number, Number],
        side_from: int,
    ) -> None:
        assert side_from in [0, 1]
        self.main_pos_0 = main_pos_0
        self.main_pos_1 = main_pos_1
        self.side_pos = side_pos
        self.side_from = side_from  # 0 for main_pos_1, 1 for main_pos_2
        self.state = 0
        self.track_object_main = Track.create(new_id())  # 正位轨道对象
        self.track_object_side = Track.create(new_id())  # 反位轨道对象

    def get_track(
        self, state: int
    ) -> Tuple[Track, Tuple[Number, Number], Tuple[Number, Number]]:
        if state == 0:
            return (self.track_object_main, self.main_pos_0, self.main_pos_1)
        else:
            if self.side_from == 0:
                return (self.track_object_side, self.main_pos_0, self.side_pos)
            else:
                return (self.track_object_side, self.side_pos, self.main_pos_1)

    def draw(self, surface: pygame.Surface) -> None:
        pygame.draw.line(
            surface,
            (255, 255, 255),
            self.main_pos_0,
            self.main_pos_1,
            4 if self.state == 0 else 2,
        )
        pygame.draw.line(
            surface,
            (255, 255, 255),
            self.main_pos_0 if self.side_from == 0 else self.main_pos_1,
            self.side_pos,
            4 if self.state == 1 else 2,
        )
        pygame.draw.circle(
            surface,
            (255, 255, 255),
            self.main_pos_0,
            3,
        )

        pygame.draw.circle(
            surface,
            (255, 255, 255),
            self.main_pos_1,
            3,
        )
        pygame.draw.circle(
            surface,
            (255, 255, 255),
            self.side_pos,
            3,
        )


class TrackModel:
    """
    轨道模型。表示一条轨道，包括起点和终点，以及轨道的状态。
    """
    def __init__(
        self, start_pos: Tuple[Number, Number], end_pos: Tuple[Number, Number]
    ) -> None:
        self.start_x, self.start_y = start_pos
        self.end_x, self.end_y = end_pos
        self.semaphore_up_state: SemaphoreStates = "PASS"
        self.semaphore_down_state: SemaphoreStates = "PASS"
        self.track_object = Track.create(new_id())

    def get_semaphore_color(
        self, semaphore: SemaphoreStates
    ) -> Tuple[Number, Number, Number]:
        if semaphore == "PASS":
            return (0, 255, 0)
        elif semaphore == "STOP":
            return (255, 0, 0)
        else:
            return (255, 255, 0)

    # def create_track(self) -> None:
    #     return Track.create()

    def draw(self, surface: pygame.Surface) -> None:

        pygame.draw.line(
            surface,
            (255, 255, 255),
            (self.start_x, self.start_y),
            (self.end_x, self.end_y),
            2,
        )
        pygame.gfxdraw.aacircle(
            surface,
            self.start_x,
            self.start_y,
            3,
            (255, 255, 255),
        )
        pygame.gfxdraw.aacircle(
            surface,
            self.end_x,
            self.end_y,
            3,
            (255, 255, 255),
        )
        draw_equilateral_triangle(
            surface,
            self.get_semaphore_color(self.semaphore_up_state),
            (self.start_x, self.start_y + 8),
            10,
            "left",
        )
        draw_equilateral_triangle(
            surface,
            self.get_semaphore_color(self.semaphore_down_state),
            (self.end_x, self.end_y + 8),
            10,
            "right",
        )
